from fastapi import APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect, status
from sqlalchemy.orm import Session
from .. import models, schemas
from ..crud import chat as chat_crud
from ..database import get_db
from ..deps import get_current_user
from typing import List, Dict
import json
from datetime import datetime

router = APIRouter()

# 存储活跃的WebSocket连接
active_connections: Dict[int, WebSocket] = {}

@router.websocket("/ws/{user_id}")
async def websocket_endpoint(websocket: WebSocket, user_id: str):
    """WebSocket连接处理"""
    await websocket.accept()
    active_connections[user_id] = websocket
    
    # 用户上线时，发送所有未读消息
    db = next(get_db())
    unread_messages = chat_crud.get_unread_messages(db, user_id)
    for message in unread_messages:
        await websocket.send_json({
            "type": "friend_message",
            "message": {
                "id": message.id,
                "content": message.content,
                "from_user_id": message.from_user_id,
                "created_at": message.created_at.isoformat()
            }
        })
    
    try:
        while True:
            data = await websocket.receive_text()
            message_data = json.loads(data)
            
            if message_data["type"] == "friend_message":
                await handle_friend_message(user_id, message_data)
            elif message_data["type"] == "avatar_message":
                await handle_avatar_message(user_id, message_data)
    except WebSocketDisconnect:
        active_connections.pop(user_id, None)

async def handle_friend_message(user_id: str, message_data: dict):
    """处理好友消息"""
    to_user_id = message_data["to_user_id"]
    content = message_data["content"]
    
    # 保存消息到数据库
    db = next(get_db())
    message = chat_crud.create_friend_message(
        db=db,
        message=schemas.FriendChatMessageCreate(
            to_user_id=to_user_id,
            content=content
        ),
        from_user_id=user_id
    )
    
    # 如果接收者在线，发送WebSocket消息
    if to_user_id in active_connections:
        await active_connections[to_user_id].send_json({
            "type": "friend_message",
            "message": {
                "id": message.id,
                "content": message.content,
                "from_user_id": message.from_user_id,
                "created_at": message.created_at.isoformat()
            }
        })

async def handle_avatar_message(user_id: str, message_data: dict):
    """处理化身消息"""
    avatar_id = message_data["avatar_id"]
    content = message_data["content"]
    
    # 保存用户消息到数据库
    db = next(get_db())
    message = chat_crud.create_avatar_message(
        db=db,
        message=schemas.AvatarChatMessageCreate(
            avatar_id=avatar_id,
            content=content
        ),
        user_id=user_id,
        is_from_user=True
    )
    
    # 调用化身AI生成回复
    from ..api.avatar_agent import get_avatar_response
    avatar = db.query(models.Avatar).filter(models.Avatar.id == avatar_id).first()
    ai_response = await get_avatar_response(avatar, content)
    
    # 保存AI回复到数据库
    ai_message = chat_crud.create_avatar_message(
        db=db,
        message=schemas.AvatarChatMessageCreate(
            avatar_id=avatar_id,
            content=ai_response
        ),
        user_id=user_id,
        is_from_user=False
    )
    
    # 发送AI回复给用户
    if user_id in active_connections:
        await active_connections[user_id].send_json({
            "type": "avatar_message",
            "message": {
                "id": ai_message.id,
                "content": ai_message.content,
                "avatar_id": avatar_id,
                "created_at": ai_message.created_at.isoformat()
            }
        })

# REST API endpoints
@router.get("/sessions", response_model=List[schemas.ChatSessionBase])
async def get_chat_sessions(
    current_user: models.User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """获取用户的所有聊天会话"""
    return chat_crud.get_user_chat_sessions(db, current_user.id)

@router.get("/friends/{friend_id}/history", response_model=List[schemas.FriendChatMessage])
async def get_friend_chat_history(
    friend_id: str,
    skip: int = 0,
    limit: int = 50,
    current_user: models.User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """获取与指定好友的聊天历史"""
    return chat_crud.get_friend_chat_history(db, current_user.id, friend_id, skip, limit)

@router.get("/avatars/{avatar_id}/history", response_model=List[schemas.AvatarChatMessage])
async def get_avatar_chat_history(
    avatar_id: str,
    skip: int = 0,
    limit: int = 50,
    current_user: models.User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """获取与指定化身的聊天历史"""
    return chat_crud.get_avatar_chat_history(db, current_user.id, avatar_id, skip, limit)

@router.post("/friends/{friend_id}/mark-read")
async def mark_friend_messages_as_read(
    friend_id: str,
    current_user: models.User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """将来自指定好友的消息标记为已读"""
    chat_crud.mark_messages_as_read(db, current_user.id, friend_id)
    return {"message": "success"}

@router.get("/unread", response_model=List[schemas.FriendChatMessage])
async def get_unread_messages(
    current_user: models.User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """获取所有未读消息"""
    return chat_crud.get_unread_messages(db, current_user.id)

@router.get("/friends/{friend_id}/unread-count")
async def get_friend_unread_count(
    friend_id: str,
    current_user: models.User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """获取来自特定好友的未读消息数量"""
    count = chat_crud.get_unread_count(db, current_user.id, friend_id)
    return {"unread_count": count} 